import numpy as np
# imports
import pytest  # used for our unit tests
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov
from scipy.spatial import ConvexHull


# Minimal RepeatedGame and StageGame classes for testing
class StageGame:
    def __init__(self, payoff_arrays):
        self.N = len(payoff_arrays)
        self.payoff_arrays = tuple(np.array(p) for p in payoff_arrays)
        self.nums_actions = tuple(p.shape[0] for p in payoff_arrays)
        # payoff_profile_array: all possible action profiles, shape (A1, A2, 2)
        # For 2 players, payoff_profile_array[i, j, :] = [p1 payoff, p2 payoff]
        A1, A2 = self.nums_actions
        self.payoff_profile_array = np.zeros((A1, A2, 2))
        for i in range(A1):
            for j in range(A2):
                self.payoff_profile_array[i, j, 0] = self.payoff_arrays[0][i, j]
                self.payoff_profile_array[i, j, 1] = self.payoff_arrays[1][j, i]

class RepeatedGame:
    def __init__(self, payoff_arrays, delta):
        self.sg = StageGame(payoff_arrays)
        self.delta = delta

# The function to test is already defined above.

# -------------- BASIC TEST CASES --------------


def test_basic_asymmetric_game():
    # Battle of the Sexes
    # Player 0: [[2, 0], [0, 1]]
    # Player 1: [[1, 0], [0, 2]]
    payoff_arrays = [
        np.array([[2, 0], [0, 1]]),
        np.array([[1, 0], [0, 2]])
    ]
    delta = 0.8
    rpg = RepeatedGame(payoff_arrays, delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 4.96ms -> 5.12ms (3.09% slower)
    points = hull.points[hull.vertices]







def test_edge_non_two_player_game():
    # Should raise NotImplementedError for N != 2
    payoff_arrays = [
        np.ones((2, 2, 2)),
        np.ones((2, 2, 2)),
        np.ones((2, 2, 2))
    ]
    class StageGame3:
        def __init__(self, payoff_arrays):
            self.N = 3
            self.payoff_arrays = tuple(np.array(p) for p in payoff_arrays)
            self.nums_actions = tuple(p.shape[0] for p in payoff_arrays)
            self.payoff_profile_array = np.ones((2, 2, 2, 3))
    class RepeatedGame3:
        def __init__(self, payoff_arrays, delta):
            self.sg = StageGame3(payoff_arrays)
            self.delta = delta
    rpg = RepeatedGame3(payoff_arrays, 0.5)
    with pytest.raises(NotImplementedError):
        _equilibrium_payoffs_abreu_sannikov(rpg) # 1.58μs -> 1.46μs (8.50% faster)

def test_edge_zero_actions():
    # Zero actions for one player (invalid)
    payoff_arrays = [
        np.zeros((0, 2)),
        np.zeros((2, 0))
    ]
    delta = 0.9
    rpg = RepeatedGame(payoff_arrays, delta)
    # Should raise ValueError or similar due to empty arrays
    with pytest.raises(ValueError):
        _equilibrium_payoffs_abreu_sannikov(rpg) # 16.4μs -> 16.6μs (1.10% slower)

def test_edge_large_tol():
    # Large tolerance should cause early stopping and possibly inaccurate result
    payoff_arrays = [
        np.array([[1, 2], [3, 4]]),
        np.array([[4, 3], [2, 1]])
    ]
    delta = 0.95
    rpg = RepeatedGame(payoff_arrays, delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg, tol=1); hull = codeflash_output # 228μs -> 212μs (7.77% faster)
    points = hull.points[hull.vertices]
    # Should still contain at least one of the input payoff pairs
    expected = np.array([[1, 4], [2, 3], [3, 2], [4, 1]])

def test_edge_max_iter_limit():
    # Setting max_iter to a low value should cause early termination
    payoff_arrays = [
        np.array([[1, 2], [3, 4]]),
        np.array([[4, 3], [2, 1]])
    ]
    delta = 0.95
    rpg = RepeatedGame(payoff_arrays, delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg, max_iter=1); hull = codeflash_output # 181μs -> 172μs (5.13% faster)
    # Should contain at least one of the input payoff pairs
    expected = np.array([[1, 4], [2, 3], [3, 2], [4, 1]])
    points = hull.points[hull.vertices]

# -------------- LARGE SCALE TEST CASES --------------

def test_large_scale_10x10_uniform_payoffs():
    # 10 actions per player, payoffs are uniform random in [0,1]
    np.random.seed(42)
    A = 10
    payoff_arrays = [
        np.random.rand(A, A),
        np.random.rand(A, A)
    ]
    delta = 0.99
    rpg = RepeatedGame(payoff_arrays, delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 62.1ms -> 56.5ms (9.85% faster)
    points = hull.points[hull.vertices]
    # Each vertex should be within the range [0,1] for both players
    for p in points:
        pass



def test_large_scale_max_actions_limit():
    # Test with the upper limit of 50 actions per player (50x50)
    np.random.seed(2024)
    A = 50
    payoff_arrays = [
        np.random.randint(-10, 10, size=(A, A)),
        np.random.randint(-10, 10, size=(A, A))
    ]
    delta = 0.8
    rpg = RepeatedGame(payoff_arrays, delta)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 169ms -> 156ms (7.82% faster)
    points = hull.points[hull.vertices]
    # Vertices should be within [-10,10] for both players
    for p in points:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import numpy as np
# imports
import pytest  # used for our unit tests
from quantecon.game_theory.repeated_game import \
    _equilibrium_payoffs_abreu_sannikov
from scipy.spatial import ConvexHull

# --- Minimal RepeatedGame and StageGame classes for testing ---

class StageGame:
    """
    Minimal stage game for two players.
    """
    def __init__(self, payoff_arrays):
        self.payoff_arrays = payoff_arrays  # tuple of 2D arrays, one per player
        self.N = 2
        self.nums_actions = tuple(arr.shape[0] for arr in payoff_arrays)
        # Payoff profile array: flatten all action profiles to shape (num_profiles, 2)
        self.payoff_profile_array = np.column_stack([
            arr.flatten() for arr in payoff_arrays
        ])

class RepeatedGame:
    """
    Minimal repeated game for two players.
    """
    def __init__(self, payoff_arrays, delta):
        self.sg = StageGame(payoff_arrays)
        self.delta = delta

# --- Unit tests for _equilibrium_payoffs_abreu_sannikov ---

# 1. Basic Test Cases









def test_edge_not_two_players():
    """
    Test with a 3-player game (should raise NotImplementedError).
    """
    class StageGame3:
        def __init__(self):
            self.N = 3
            self.nums_actions = (2, 2, 2)
            arr = np.zeros((2, 2, 2))
            self.payoff_arrays = (arr, arr, arr)
            self.payoff_profile_array = np.zeros((8, 3))
    class RepeatedGame3:
        def __init__(self):
            self.sg = StageGame3()
            self.delta = 0.9
    rpg = RepeatedGame3()
    with pytest.raises(NotImplementedError):
        _equilibrium_payoffs_abreu_sannikov(rpg) # 1.56μs -> 1.49μs (4.84% faster)



def test_large_scale_10x10_payoff():
    """
    Test with a 10x10 game, random payoffs in [0, 10], delta=0.9.
    """
    np.random.seed(42)
    payoff0 = np.random.randint(0, 11, size=(10, 10))
    payoff1 = np.random.randint(0, 11, size=(10, 10))
    rpg = RepeatedGame((payoff0, payoff1), delta=0.9)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 35.7ms -> 32.5ms (9.62% faster)
    # Hull should contain points from the payoff arrays
    expected_points = np.column_stack([payoff0.flatten(), payoff1.flatten()])
    # Pick 5 random points to check
    for i in np.random.choice(expected_points.shape[0], 5, replace=False):
        pt = expected_points[i]

def test_large_scale_30x30_payoff():
    """
    Test with a 30x30 game, random payoffs in [-20, 20], delta=0.95.
    """
    np.random.seed(123)
    payoff0 = np.random.randint(-20, 21, size=(30, 30))
    payoff1 = np.random.randint(-20, 21, size=(30, 30))
    rpg = RepeatedGame((payoff0, payoff1), delta=0.95)
    codeflash_output = _equilibrium_payoffs_abreu_sannikov(rpg); hull = codeflash_output # 393ms -> 335ms (17.2% faster)
    # Hull should contain at least one point from payoff arrays
    expected_points = np.column_stack([payoff0.flatten(), payoff1.flatten()])
    pt = expected_points[0]